'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-11-19
Copyright © Department of Physics, Tsinghua University. All rights reserved

Converting tags to directed acyclic graph.
Raw data is in data/tags.txt. Format:
    tag1
    tag2
    tag3
    tag1/tag2
    tag1/tag3
    tag2/tag3
    tag4/tag5/tag6
    tag7/tag8/tag9

Result:
    tag1 -> tag2
    tag1 -> tag3
    tag2 -> tag3
    tag4 -> tag5 -> tag6
    tag7 -> tag8 -> tag9
'''

import networkx
import literature_format as lf


def tag_str_to_edges(tag_str):
    ''' tag1/tag2/tag3 -> [(tag1, tag2), (tag2, tag3)] '''
    # processing tags
    tag_list = tag_str.split('/')
    for i in range(len(tag_list) - 1):
        yield tag_list[i], tag_list[i + 1]


def expand_tags(tag_list: list[str]) -> set[str]:
    expanded_tags = set(tag_list)  # Start with the original tags in a set

    for tag in tag_list:
        parts = tag.split('/')
        for i in range(1, len(parts)):
            expanded_tags.add('/'.join(parts[:i]))

    return expanded_tags


def check_conflicted_names(tag_list: list[str]):
    ''' Check if there are conflicted tag names. That is, the same tag name appears in different places.
        If there are, find the list of duplicated tags. Else, return empty list.
        
        Example:
            ========
            input:
                ["tag1/tag2/tag3", "tag2"]
            output:
                ["tag1/tag2", "tag2"]
            explanation:
                tag2 appears in both "tag1/tag2/tag3" and "tag2". 
                "tag2" is different from "tag1/tag2", but uses the same name.
            ========
            input:
                ["tag1/tag2/tag3", "tag1/tag2"]
            output:
                []
            explanation:
                although tag1/tag2 appears in both "tag1/tag2/tag3" and "tag1/tag2",
                the latter refers to the name of category of tag3, so that it is not a duplicated tag
            ========
            input:
                ["tag1/tag2/tag3", "tag4/tag2/tag5"]
            output:
                ["tag1/tag2", "tag4/tag2"]
            explanation:
                tag2 appears in both "tag1/tag2" and "tag4/tag2".
            ========
            input:
                ["tag1/tag2", "tag3"]
            output:
                []
            explanation:
                tag1/tag2 and tag3 are both simple tags, so they are not conflicted.
            ========
    '''

    tag_list = [tag_str.replace('#', '') for tag_str in tag_list]

    # Expand the tags
    expanded_tags = expand_tags(tag_list)
    
    # Check for conflicted names
    conflicted_names = []
    for tag in expanded_tags:
        parts = tag.split('/')
        last_part = parts[-1]
        
        # Check if the last part of the tag appears in more than one tag
        if sum(1 for t in expanded_tags if last_part == t.split('/')[-1]) > 1:
            conflicted_names.append(tag)
    
    return conflicted_names


def check_conflicted_tags_in_my_library():
    ''' Check if there are conflicted tags in my library. That is, the same tag appears in different places.
        If there are, find the list of duplicated tags. Else, return empty list.'''
    tag_list = lf.zotero.get_tags_in_library(
        lf.zotero.MY_LIBRARY_ID
    )
    return check_conflicted_names(tag_list)


def tags_to_graph(tag_list: list[str]):
    ''' Convert list of tags to DAG. If a tag is a simple tag (without '/'), it will be added to the graph. 
        If a tag is a complex tag (with '/'), it will be added to the graph as an edge. '''
    assert(check_conflicted_names(tag_list) == [])

    G = networkx.DiGraph()
    for tag_str in tag_list:
        tag_str = tag_str.replace('#', '')
        tag_str = tag_str.strip()
        if '/' in tag_str:
            G.add_edges_from(tag_str_to_edges(tag_str))
        else:
            G.add_node(tag_str)
    return G


def generate_tag_graph():
    ''' Generate tag graph from tags.txt. '''
    tag_list = lf.zotero.get_tags_in_library(
        lf.zotero.MY_LIBRARY_ID
    )
    G = tags_to_graph(tag_list)
    with open('data/tag-graph.gml', 'wb') as f:
        networkx.write_gml(G, f)


def write_tag_graph_to_text(G: networkx.DiGraph):
    ''' Write tag graph to text. 
        Format:
        ===
        keyword: "tag1"
        parent keywords: [
            "tag2",
            "tag3"
        ]
        ===
        keyword: "tag2"
        ===
    '''
    # BFS
    with open('data/tag-graph.txt', 'w', encoding='utf-8') as f:
        for node in networkx.topological_sort(G):
            f.write('keyword: "{}"\n'.format(node))
            n_predecessors = 0
            for _ in G.predecessors(node):
                n_predecessors += 1
                break
            if n_predecessors:
                f.write('parent keywords: [')
                for child in G.predecessors(node):
                    f.write('"{}",'.format(child))
                f.write(']\n')
            f.write("===\n")


def get_graph_from_my_library():
    ''' Get tag graph from my library. '''
    tag_list = lf.zotero.get_tags_in_library(
        lf.zotero.MY_LIBRARY_ID
    )

    # remove tags in tag-config/tagignore.txt
    with open('tag-config/tagignore.txt', 'r', encoding='utf-8') as f:
        tagignore_list = f.readlines()
    tagignore_list = [tagignore.strip() for tagignore in tagignore_list]
    tag_list = [tag for tag in tag_list if tag not in tagignore_list]

    # check conflicted names
    conflicted_names = check_conflicted_names(tag_list)
    assert(conflicted_names == [])

    # generate graph
    G = tags_to_graph(tag_list)
    return G
 

def write_all_tags_in_my_library():
    ''' Write all tags in my library to text. '''
    # write to text
    G = get_graph_from_my_library()
    write_tag_graph_to_text(G)


def get_all_first_level_tags():
    ''' Get all first level tags in my library. '''
    G = get_graph_from_my_library()
    with open("data/first-level-tags.txt", "w", encoding="utf-8") as fp:
        for node in G.nodes():
            if G.in_degree(node) == 0:
                print(node, file=fp)


if __name__ == "__main__":
    # check_conflicted_tags_in_my_library()
    # generate_tag_graph()
    # write_all_tags_in_my_library()
    get_all_first_level_tags()